1
スレッドからプログラムインスタンスへの移行
AI023Lesson 4
00:00

Tritonでは、実行の基本単位がCUDAスカラー スレッドから プログラムインスタンスに移行します。これは、1つのインスタンスが同時にベクトル化された「ブロック」の要素を処理するGPUスレッドブロックの抽象表現です。

1. プログラムインスタンスの識別子

すべての実行ユニットは、 pid = tl.program_id(axis=0)によって自身の識別子を取得します。 倉庫のフォークリフト (プログラムインスタンス)が パレット 128個の箱の束(ブロック)を持ち上げるのと、1人の作業者(CUDAスレッド)が1箱ずつ持ち上げるのを比較してください。

2. TritonとPyTorchテンソルの比較

メモリ管理において、この意味的なギャップを理解することは重要です:

  • PyTorchテンソル: ホスト側のPythonオブジェクトで、VRAMのストレージ、ストライド、メタデータをラップしています。
  • Tritonテンソル: コンパイラレベルのオブジェクトで、 レジスタまたはSRAMに格納されている値やポインタを表します
PyTorchビュー
連続したグローバルメモリを指すPythonオブジェクト。
Tritonビュー
コンパイラのレジスタ内の2次元/1次元のデータブロック。

3. SPMDの性質

Tritonは 単一プログラム・複数データ(SPMD) のフローに従います。すべてのプログラムインスタンスは まったく同じ コードを実行します。分岐は、論理が pid を使って特定のメモリオフセットを計算する場合にのみ発生します。

main.py
TERMINALbash — 80x24
> Ready. Click "Run" to execute.
>